import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.patches import Rectangle
import os


# --- 1. 指标计算函数 ---
# --- 添加这两行来解决中文显示问题 ---
plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
plt.rcParams['axes.unicode_minus'] = False   # 解决保存图像是负号'-'显示为方块的问题
def lwma(series: pd.Series, period: int) -> pd.Series:
    """
    计算线性加权移动平均 (Linear Weighted Moving Average - LWMA)。
    近期价格的权重最高。
    """
    weights = np.arange(1, period + 1)
    sum_weights = np.sum(weights)
    # rolling().apply() 会将窗口内的数据（作为numpy数组）传递给lambda函数
    return series.rolling(window=period).apply(lambda x: np.dot(x, weights) / sum_weights, raw=True)


def calculate_indicators(df: pd.DataFrame, period: int = 20) -> pd.DataFrame:
    """
    根据猜测的逻辑计算所有指标和交易信号。
    """
    print("正在计算指标...")

    # 一、构建趋势线 (HHD 和 LLD)
    # 这是对高点和低点的20周期线性加权移动平均
    df['HHD'] = lwma(df['high'], period)
    df['LLD'] = lwma(df['low'], period)

    # 定义通道中线，用于判断“相对高/低位”
    df['midpoint'] = (df['HHD'] + df['LLD']) / 2

    # 二、定义三种状态
    df['state'] = np.where(
        df['close'] > df['HHD'], 'up',  # 收盘价在通道上方 -> 上涨
        np.where(df['close'] < df['LLD'], 'down',  # 收盘价在通道下方 -> 下跌
                 'sideways')  # 在通道内部 -> 震荡
    )

    # 三、识别震荡区入场点 (cond1 和 cond2)
    # 计算N周期内的最高价(HH)和最低价(LL)
    df['HH'] = df['high'].rolling(window=period).max()
    df['LL'] = df['low'].rolling(window=period).min()

    # 为了应用[1]（前一周期），我们使用shift(1)
    prev_high = df['high'].shift(1)
    prev_low = df['low'].shift(1)
    prev_hh = df['HH'].shift(1)
    prev_ll = df['LL'].shift(1)
    prev_hhd = df['HHD'].shift(1)
    prev_lld = df['LLD'].shift(1)
    prev_midpoint = df['midpoint'].shift(1)
    prev_is_sideways = (df['state'].shift(1) == 'sideways')

    # cond1: 震荡区相对高位开空条件
    df['short_signal'] = (
            (prev_hh > prev_midpoint) &  # 阶段高点在中线之上
            (prev_high == prev_hh) &  # 前一根K线创下阶段新高
            (prev_is_sideways) &  # 前一根K线处于震荡状态
            (prev_high < prev_hhd)  # 但这个新高未能突破通道上轨
    )

    # cond2: 震荡区相对低位开多条件
    df['long_signal'] = (
            (prev_ll < prev_midpoint) &  # 阶段低点在中线之下
            (prev_low == prev_ll) &  # 前一根K线创下阶段新低
            (prev_is_sideways) &  # 前一根K线处于震荡状态
            (prev_low > prev_lld)  # 但这个新低未能跌破通道下轨
    )

    return df


# --- 2. 数据加载函数 ---

def load_data(filepath: str) -> pd.DataFrame:
    """
    加载并预处理CSV数据。
    """
    print(f"正在从 {filepath} 加载数据...")
    if not os.path.exists(filepath):
        raise FileNotFoundError(f"错误：文件不存在于路径 {filepath}")

    column_names = [
        'open_time', 'open', 'high', 'low', 'close', 'volume',
        'close_time', 'quote_volume', 'count', 'taker_buy_volume',
        'taker_buy_quote_volume', 'ignore'
    ]
    df = pd.read_csv(filepath, header=0, names=column_names)

    # 将时间戳转换为datetime对象，并设为索引
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)

    # 确保数据类型正确
    for col in ['open', 'high', 'low', 'close']:
        df[col] = pd.to_numeric(df[col])

    return df


# --- 3. 可视化函数 ---

def plot_chart(df: pd.DataFrame):
    """
    绘制K线图、通道、状态背景和交易信号。
    """
    print("正在生成图表...")

    fig, ax = plt.subplots(figsize=(18, 9))

    # 绘制状态背景
    for i in range(len(df)):
        state = df['state'].iloc[i]
        color_map = {'up': 'green', 'down': 'red', 'sideways': 'yellow'}
        color = color_map.get(state, 'white')

        # 使用axvspan为每个K线周期添加垂直背景色块
        start_date = df.index[i]
        # 假设是1小时数据
        end_date = start_date + pd.Timedelta(hours=1)
        ax.axvspan(start_date, end_date, facecolor=color, alpha=0.2)

    # 绘制K线图
    # 筛选出上涨和下跌的K线，分别用绿色和红色绘制
    up = df[df.close >= df.open]
    down = df[df.close < df.open]

    # 上涨K线 (绿色)
    ax.bar(up.index, up.close - up.open, 0.03, bottom=up.open, color='green', edgecolor='black')
    ax.vlines(up.index, up.low, up.high, color='green', linewidth=0.8)

    # 下跌K线 (红色)
    ax.bar(down.index, down.close - down.open, 0.03, bottom=down.open, color='red', edgecolor='black')
    ax.vlines(down.index, down.low, down.high, color='red', linewidth=0.8)

    # 绘制HHD和LLD通道线
    ax.plot(df.index, df['HHD'], color='blue', linestyle='--', label='HHD (上轨)', linewidth=1.5)
    ax.plot(df.index, df['LLD'], color='purple', linestyle='--', label='LLD (下轨)', linewidth=1.5)

    # 标记入场点
    long_signals = df[df['long_signal']]
    short_signals = df[df['short_signal']]

    ax.scatter(long_signals.index, long_signals['low'] * 0.995, marker='^', color='lime', s=100, label='做多信号',
               zorder=5)
    ax.scatter(short_signals.index, short_signals['high'] * 1.005, marker='v', color='magenta', s=100, label='做空信号',
               zorder=5)

    # 格式化图表
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d %H:%M'))
    ax.xaxis.set_major_locator(mdates.DayLocator(interval=2))  # 每两天显示一个日期
    plt.xticks(rotation=45)
    plt.title('BNB/USDT 1小时图 - 趋势线分析', fontsize=16)
    plt.ylabel('价格 (USDT)', fontsize=12)
    plt.xlabel('日期和时间', fontsize=12)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.tight_layout()

    print("图表生成完毕。")
    plt.show()


# --- 主程序入口 ---
if __name__ == "__main__":
    # 请确保将此路径修改为您电脑上的实际文件路径
    # 在Windows上，路径前的 'r' 很重要，它可以防止反斜杠被错误地解析
    filepath = r"F:\personal\binance_klines\BNBUSDT\1h\BNBUSDT-1h-2025-03.csv"

    try:
        # 1. 加载数据
        kline_data = load_data(filepath)

        # 2. 计算指标
        data_with_indicators = calculate_indicators(kline_data)

        # 3. 绘制图表 (为看得更清楚，可以只绘制一部分数据，例如后300根K线)
        plot_chart(data_with_indicators.tail(300))

    except FileNotFoundError as e:
        print(e)
        print("请检查文件路径是否正确。")
    except Exception as e:
        print(f"程序运行中发生未知错误: {e}")